"""
https://leetcode.cn/problems/diagonal-traverse/description/

498. 对角线遍历
中等
相关标签
相关企业
给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

 

示例 1：


输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,2,4,7,5,3,6,8,9]
示例 2：

输入：mat = [[1,2],[3,4]]
输出：[1,2,3,4]
 

提示：

m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
-105 <= mat[i][j] <= 105
"""

from typing import List


class Solution:




    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        """
        本质还是找规律，但原规律太丑了
        边界调节：
        0<=col<=height-1
        0<=row<=width-1
        row+col逐次增大
        当row+col奇数是，方向为ld;否则方向为ru

        todo 待提升查询效率
        原始过滤条件为：0<=c<=h-1 and 0<=sum-c<=w-1
        后者可推导出：-sum<=-c<=w-1-sum => sum+1-w<=c<=sum
        合计则有：max(0,sum+1-w)<=c<=min(h-1,sum)

        """
        h=len(mat)
        w=len(mat[0])
        
        res_list = []
        sr=c=0
        for sum in range(w+h-1):
            # print(f'sum={sum}')
            is_odd=sum%2==1
            min_row=max(0,sum+1-w)
            max_row=min(h-1,sum)
            sr=min_row if is_odd else max_row
            er=max_row if is_odd else min_row
            step=1 if is_odd else -1
            for r in range(sr,er+step,step):
                res_list.append(mat[r][sum-r])
                # if 0<=c<=h-1 and 0<=sum-c<=w-1:          
        
        return res_list



if __name__=='__main__':
    sol=Solution()
    mat = [[2,3]]
    res=sol.findDiagonalOrder(mat)
    print(res)